TODO: Improve this. For sure try to capture HTTP errors in more verbose logs to avoid excessive log generation if the error is for instance server raising 500 status code.

I guess to simplify this, we could just pass an instance of the exception here instead of the class and message? That would additionally also allow some exceptions which can/need to be instantiated with multiple arguments.

This should probably be of type ExecutionStatus?

My dynamic pipelines PR introduces running multiple steps in different threads, which doesn't work with this I think.

Can we somehow store the thread from which the heartbeat worker was started, and then interrupt that thread instead of the main one?

Yeah that is an important change, good point. interrupt_main will not work here, we will need to change the pattern a bit. Should I work my changes from your branch?

I'm wondering whether this RBAC check is even necessary, as running all of this will take quite some time (two calls to the DB, then a request to the RBAC service).

Is there any real harm in leaving this unprotected? I guess it would allow users potential access to the status of the step, which I'm not sure really is a concern.

True, we can probably do both authenication & authorization with pipeline tokens. Will discuss with @stefannica for directions.

Let me suggest an alternative: we could limit this endpoint to only be accessed by running pipelines.

Running pipelines (the containerized environment where the steps are running actually) use something called "a workload API token" which is only valid as long as the pipeline run itself is not yet finalized. These workload API tokens are tied to a particular pipeline run (or schedule, in case of scheduled pipelines). So we can also use their scope to limit the range of targets that they can update.

Some references:

• this is the code that verifies the pipeline scoped tokens (you can see some leeway is involved): https://github.com/zenml-io/zenml/blob/main/src/zenml/zen_server/auth.py#L406-L475
• same thing for the schedule-scoped tokens: https://github.com/zenml-io/zenml/blob/main/src/zenml/zen_server/auth.py#L363-L404

A sketch of how you can use this in your endpoint:

def update_heartbeat(
    step_run_id: UUID,
    auth_context: AuthContext = Security(authorize),
) -> StepHeartbeatResponse:

    ...
    if not auth_context.access_token or not auth_context.access_token.schedule_id and not auth_context.access_token.pipeline_run_id:
        raise AuthorizationException("Not authorized")

    if auth_context.access_token.pipeline_run_id:
        # optionally, check that the step ID is part of this run ID
    else: # if auth_context.access_token.schedule_id
        # optionally, check that the step ID is part of a run ID that was scheduled with this schedule

This will no longer rely on RBAC calls, but it might still flood the database with a lot of requests, so maybe you could also implement a mini-caching system like the ones used in the previous code references, to reduce its impact.

@stefannica That was my initial idea as well, but do we use those tokens also when running pipelines with service accounts? I thought at some point we used the API key directly when running scheduled pipelines, but I might be misremembering.

I know for sure though that there is a way to generate a generic unscoped token instead of a workload token when running a pipeline (by setting some token expiration env variable), so we'll have to think about how we handle this case.

@schustmi yes, even when running pipelines with service accounts, we generate workload API tokens scoped to the pipeline run or schedule. The only case where we use a generic unscoped token is if you set the ZENML_PIPELINE_API_TOKEN_EXPIRATION env variable. But this is a very obscure case, which I don't think we need to handle separately. In that case, we can just not run any RBAC like checks on this endpoint.

The problem is that this is a client-side env variable and I'm not sure we can recognize this case in the server endpoint that receives the heartbeat requests. It will just be a generic token that is not scoped to the run, and if we allow those to call the endpoint without any checks then everyone can do so, no?

Maybe instead of this we can check pipeline_run.schedule.id == ctx.access_token.schedule_id, to make sure the run is actually triggered by the schedule that the token is scoped to?

Nice, missed that. Yes, that's better 1 DB call shorter.

Just a question for clarification: If someone actually interrupts the python process with CTRL-C, that raises a KeyboardInterrupt in the main thread I assume?

It would be captured the same way and a misleading exception would appear I would think. But is that a scenario we need to worry about?

I guess I personally do it quite a lot which is why I thought about it, not sure how common/important it is though. I guess it's quite easy to solve though, so maybe we can account for it?

It could just be a boolean flag on the heartbeat worker that signals that it interrupted the main thread, and we can use that flag to decide whether to re-raise with the HeartbeatInterrupt exception or keep the original KeyboardInterrupt?

Yeap, that flag already exists in the heartbeat it is .is_terminated. Easy to do, will cover that as well.

I'm guessing you don't stop the worker here and instead rely on the server response for the thread to shut down automatically?

Wouldn't it be better to add an explicit stop once the user code finished executing?

Yeap makes sense to add it. Won't make much of a difference now but it is safer and more future-proof.